#include "CUDA_DTIFiberTracking.h"

#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"

/* DEFINES */
#define CUDA_TEN_T2M(m, t) ( \
   (m)[0] = (t)[1], (m)[1] = (t)[2], (m)[2] = (t)[3], \
   (m)[3] = (t)[2], (m)[4] = (t)[4], (m)[5] = (t)[5], \
   (m)[6] = (t)[3], (m)[7] = (t)[5], (m)[8] = (t)[6] )

__shared__ float3 sample_points[SAMPLE_POINT_NUM];
__shared__ float sample_point_weights[SAMPLE_POINT_NUM];

__shared__ float stepSize;
__shared__ float maxLength;
__shared__ float faThreshold;
__shared__ float angleThreshold;
__shared__ float MLSSize;
__shared__ int innerSteps;

/* textures */
/* MASK */
texture<float, cudaTextureType3D> texMASK;
/* DTI */
texture<float4, cudaTextureType3D> texDTI_1;
texture<float4, cudaTextureType3D> texDTI_2;


__device__ float CUDA_DOT(const float3 *op1, const float3 *op2)
{
	float res = (op1->x) * (op2->x) + (op1->y) * (op2->y) + (op1->z) * (op2->z);
	return res;
}

__device__ float CUDA_GetMASK(const float3 *pos)
{
	float x = pos->x + 0.5f;
	float y = pos->y + 0.5f;
	float z = pos->z + 0.5f;
	float mask = tex3D(texMASK, x, y, z);
	return mask;
}

__device__ void CUDA_GetDTI(const float3 *pos, float *dti)
{
	float x = pos->x + 0.5f;
	float y = pos->y + 0.5f;
	float z = pos->z + 0.5f;

	float4 temp;

	temp = tex3D(texDTI_1, x, y, z);
	dti[0] = temp.x;
	dti[1] = temp.y;
	dti[2] = temp.z;
	dti[3] = temp.w;

	temp = tex3D(texDTI_2, x, y, z);
	dti[4] = temp.x;
	dti[5] = temp.y;
	dti[6] = temp.z;
}


__device__ float CUDA_DTI2FA(const float *dti)
{
	float t[6];
	t[0] = dti[1]; t[1] = dti[2]; t[2] = dti[3]; t[3] = dti[4]; t[4] = dti[5]; t[5] = dti[6];

	float mean = (t[0] + t[1] + t[2] + t[3] + t[4] + t[5]) / 6.0f;
	if (mean == 0.0f)
		return 0.0f;

	t[0] = t[0] / mean;
	t[1] = t[1] / mean;
	t[2] = t[2] / mean;
	t[3] = t[3] / mean;
	t[4] = t[4] / mean;
	t[5] = t[5] / mean;

	float cross = t[1] * t[1] + t[2] * t[2] + t[4] * t[4];
	float j2 = t[0] * t[3] + t[3] * t[5] + t[5] * t[0] - cross;
	float j4 = t[0] * t[0] + t[3] * t[3] + t[5] * t[5] + 2.0f * cross;
	float fa = sqrtf((j4-j2)/j4);
	return fa;
}

__device__ int CUDA_DTI2Orientation(const float *dti, float3 *orientation)
{
	float t[6];
	t[0] = dti[1]; t[1] = dti[2]; t[2] = dti[3]; t[3] = dti[4]; t[4] = dti[5]; t[5] = dti[6];

	float mean = (t[0] + t[1] + t[2] + t[3] + t[4] + t[5]) / 6.0f;
	t[0] = t[0] / mean;
	t[1] = t[1] / mean;
	t[2] = t[2] / mean;
	t[3] = t[3] / mean;
	t[4] = t[4] / mean;
	t[5] = t[5] / mean;
	// comput the major eigen value
	float I1 = t[0] + t[3] + t[5];
	float I2 = t[0] * t[3] + t[0] * t[5] + t[3] * t[5] - 
		(t[1] * t[1] + t[2] * t[2] + t[4] * t[4]);
	float I3 = t[0] * t[3] * t[5] + 2.0f * (t[1] * t[2] * t[4]) - 
		(t[0] * t[4] * t[4] + t[3] * t[2] * t[2] + t[5] * t[1] * t[1]);

	float third = 1.0f / 3.0f;
	float I1third = I1 * third;
	float I1thirdsqr = I1third * I1third;
	float I2third = I2 * third;
	float v = I1thirdsqr - I2third;

	float s = I1thirdsqr * I1third - I1 * I2 / 6.0f + 0.5f * I3;

	if (v < 0.0f)
		v = 0.0f;
	float sqrtv = sqrtf(v);
	float temp = s / (v * sqrtv);
	if (temp > 1.0f)
		temp = 1.0f;
	if (temp < -1.0f)
		temp = -1.0f;
	float phi = acosf(temp) * third;
	float sqrtv2 = 2.0f * sqrtv;

	float e1 = I1third + sqrtv2 * cosf(phi);

	// compute the major eigen vector
	float3 v1;
	v1.x = (t[1] * t[4] - (t[3] - e1) * t[2]) * (t[2] * t[4] - (t[5] - e1) * t[1]);
	v1.y = (t[4] * t[2] - (t[5] - e1) * t[1]) * (t[1] * t[2] - (t[0] - e1) * t[4]);
	v1.z = (t[4] * t[1] - (t[3] - e1) * t[2]) * (t[2] * t[1] - (t[0] - e1) * t[4]);

	// normalize
	float len = sqrtf(v1.x * v1.x + v1.y * v1.y + v1.z * v1.z);
	if (len > 0.0f) {
		v1.x = v1.x / len;
		v1.y = v1.y / len;
		v1.z = v1.z / len;
	}

	*orientation = v1;

	return 1;
}


__device__ void CUDA_DTI_Filter(const float3 *pos, float *dti)
{
	float temp_dti[7];
	CUDA_GetDTI(pos, temp_dti);

	float m[9];
	CUDA_TEN_T2M(m, temp_dti);

	/* get the maximum length */
	float max_length = 0.0f;
	for (int i = 0; i < SAMPLE_POINT_NUM; ++i) {
		float3 sample;
		sample.x = sample_points[i].x * m[0] + sample_points[i].y * m[1] + sample_points[i].z * m[2];
		sample.y = sample_points[i].x * m[3] + sample_points[i].y * m[4] + sample_points[i].z * m[5];
		sample.z = sample_points[i].x * m[6] + sample_points[i].y * m[7] + sample_points[i].z * m[8];
		float length = sqrtf(CUDA_DOT(&sample, &sample));
		if (length > max_length)
			max_length = length;
	}

	if (max_length < 10e-7)
		max_length = 1.0f;

	/* filter */
	for (int i = 0; i < 7; ++i)
		dti[i] = 0.0f;
	float weight = 0.0f;
	for (int i = 0; i < SAMPLE_POINT_NUM; ++i) {
		float3 sample;
		sample.x = sample_points[i].x * m[0] + sample_points[i].y * m[1] + sample_points[i].z * m[2];
		sample.y = sample_points[i].x * m[3] + sample_points[i].y * m[4] + sample_points[i].z * m[5];
		sample.z = sample_points[i].x * m[6] + sample_points[i].y * m[7] + sample_points[i].z * m[8];
		sample.x = sample.x / max_length * MLSSize + pos->x;
		sample.y = sample.y / max_length * MLSSize + pos->y;
		sample.z = sample.z / max_length * MLSSize + pos->z;
		if (CUDA_GetMASK(&sample) < 0.1f) {
			continue;
		}
		CUDA_GetDTI(&sample, temp_dti);
		for (int j = 0; j < 7; ++j) {
			dti[j] += sample_point_weights[i] * temp_dti[j];
		}
		weight += sample_point_weights[i];
	}
	
	if (weight < 10e-7)
		weight = 1.0f;

	for (int j = 0; j < 7; ++j) {
		dti[j] = dti[j] / weight;
	}
}

__device__ int CUDA_GetOrientation(const float3 *in_pos, const float3 *in_dir, float3 *out_dir)
{
	out_dir->x = out_dir->y = out_dir->z = 0.0f;

	float mask = CUDA_GetMASK(in_pos);
	if (mask < 0.1f)
		return 0;

	float dti[7];
	CUDA_DTI_Filter(in_pos, dti);

	float dti_sum = fabs(dti[1]) + fabs(dti[2]) + fabs(dti[3]) + fabs(dti[4]) + fabs(dti[5]) + fabs(dti[6]);
	if (dti_sum < 10e-7)
		return 0;

	float fa = CUDA_DTI2FA(dti);
	if (fa < faThreshold)
		return 0;

	CUDA_DTI2Orientation(dti, out_dir);

	if (in_dir == NULL) {                       /* no input orientations */
		return 1;
	} else {                                    /* select the best one */
		float v = CUDA_DOT(out_dir, in_dir);
		/* check the angle threshold */
		if (fabs(v) < angleThreshold)
			return 0;

		/* re-orientation */
		if (v < 0.0f) {
			out_dir->x = -(out_dir->x);
			out_dir->y = -(out_dir->y);
			out_dir->z = -(out_dir->z);
		}
		return 1;
	}
}

__device__ void CUDA_UpdatePosition(const float3 *last_pos, const float3 *last_dir, float3 *next_pos)
{
	next_pos->x = last_pos->x + last_dir->x * stepSize;
	next_pos->y = last_pos->y + last_dir->y * stepSize;
	next_pos->z = last_pos->z + last_dir->z * stepSize;
}

__device__ int CUDA_ComputeNextPositionRK2(const float3 *last_pos, const float3 *last_dir,
										   float3 *next_pos, float3* next_dir)
{
	float3 p;
	p.x = last_pos->x + last_dir->x * stepSize / 2.0f;
	p.y = last_pos->y + last_dir->y * stepSize / 2.0f;
	p.z = last_pos->z + last_dir->z * stepSize / 2.0f;

	float3 v;
	if (CUDA_GetOrientation(&p, last_dir, &v) == 0) {
		return 0;
	}

	next_pos->x = last_pos->x + v.x * stepSize;
	next_pos->y = last_pos->y + v.y * stepSize;
	next_pos->z = last_pos->z + v.z * stepSize;

	if (CUDA_GetOrientation(next_pos, last_dir, next_dir) == 0) {
		return 0;
	}

	return 1;
}


__global__ void 
GenerateSeeds_kernel(const float3 *points, const float *weights, CUDA_Seed *seeds, 
					 const float in_faThreshold, const float in_MLSSize, const int count)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	int index = y * gridDim.x * blockDim.x + x;

	if (threadIdx.x == 0) {
		for (int i = 0; i < SAMPLE_POINT_NUM; ++i) {
			sample_points[i] = points[i];
			sample_point_weights[i] = weights[i];
		}

		 faThreshold = in_faThreshold;
		 MLSSize = in_MLSSize;
	}
	__syncthreads();

	if (index < count) {
		float3 pos = seeds[index].pos;

		float3 dir;
		int num = CUDA_GetOrientation(&pos, NULL, &dir);

		seeds[index].dir = dir;

		if (num == 0)
			seeds[index].id = -1;
	}

	//__syncthreads();
}



__global__ void
TrackFibers_kernel(const float3 *points, const float *weights, CUDA_Seed *seeds, 
				   CUDA_Fiber *fibers, const float in_stepSize, const float in_maxLength, 
				   const float in_faThreshold, const float in_angleThreshold, 
				   const float in_MLSSize, const int in_innerSteps, const int count)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	int index = y * gridDim.x * blockDim.x + x;

	if (threadIdx.x == 0) {
		for (int i = 0; i < SAMPLE_POINT_NUM; ++i) {
			sample_points[i] = points[i];
			sample_point_weights[i] = weights[i];
		}

		stepSize = in_stepSize;
		maxLength = in_maxLength;
		faThreshold = in_faThreshold;
		angleThreshold = in_angleThreshold;
		innerSteps = in_innerSteps;
		MLSSize = in_MLSSize;
	}
	__syncthreads();

	if (index < count) {
		const CUDA_Seed *seed = &(seeds[index]);
		CUDA_Fiber *fiber = &(fibers[index]);

		int maxSteps = (int)(maxLength / (stepSize * (float)innerSteps));

		int fnum = 0;
		int bnum = 0;
		float3 last_dir = seed->dir;

		/* tracking forward */
		fiber->f[0] = seed->pos;
		for (fnum = 1; fnum < maxSteps; ++fnum) {
			float3 last_pos = fiber->f[fnum-1];
			bool flag = false;
			for (int step = 0; step < innerSteps; ++step) {
				float3 next_pos, next_dir;
				if (CUDA_ComputeNextPositionRK2(&last_pos, &last_dir, &next_pos, &next_dir) == 0) {
					flag = true;
					break;
				}
				last_dir = next_dir;
				last_pos = next_pos;
			}
			if (flag == true)
				break;
			fiber->f[fnum] = last_pos;
		}

		///* tracking backward */
		last_dir.x = -seed->dir.x; last_dir.y = -seed->dir.y; last_dir.z = -seed->dir.z;
		fiber->b[0] = seed->pos;
		for (bnum = 1; bnum < maxSteps; ++bnum) {
			float3 last_pos = fiber->b[bnum-1];
			bool flag = false;
			for (int step = 0; step < innerSteps; ++step) {
				float3 next_pos, next_dir;
				if (CUDA_ComputeNextPositionRK2(&last_pos, &last_dir, &next_pos, &next_dir) == 0) {
					flag = true;
					break;
				}
				last_dir = next_dir;
				last_pos = next_pos;
			}
			if (flag == true)
				break;
			fiber->b[bnum] = last_pos;
		}

		fiber->fc = fnum;
		fiber->bc = bnum;

		fiber->seed.dir = seed->dir;
		fiber->seed.pos = seed->pos;
		fiber->seed.id = seed->id;
	}
}


/************************************************************************/
/************************************************************************/
/************************************************************************/
extern "C"
void CUDA_GenerateSeeds(GPUPlan *plan, const float in_faThreshold, 
						const float in_MLSSize, const int num)
{
	cutilSafeCall(cudaSetDevice(plan->deviceID));
	cutilSafeCall(cudaMemcpyAsync(plan->d_pSeeds, plan->h_pSeeds, 
		THREADS_PER_RUN*sizeof(CUDA_Seed), cudaMemcpyHostToDevice, plan->stream));

	dim3 dimBlock(64, 1);
	dim3 dimGrid(THREADS_PER_RUN/64, 1);
	GenerateSeeds_kernel<<<dimGrid, dimBlock, 0, plan->stream>>>
		(plan->d_pSamplePoints, plan->d_pSampleWeights, plan->d_pSeeds, in_faThreshold, in_MLSSize, num);

	cutilSafeCall(cudaGetLastError());

	cutilSafeCall(cudaMemcpyAsync(plan->h_pSeeds, plan->d_pSeeds, 
		THREADS_PER_RUN*sizeof(CUDA_Seed), cudaMemcpyDeviceToHost, plan->stream));
}

extern "C"
void CUDA_TrackFibers(GPUPlan *plan, const float in_stepSize, const float in_maxLength, 
					  const float in_faThreshold, const float in_angleThreshold, 
					  const float in_MLSSize, const int in_innerSteps, const int num)
{
	cutilSafeCall(cudaSetDevice(plan->deviceID));
	cutilSafeCall(cudaMemcpyAsync(plan->d_pSeeds, plan->h_pSeeds, 
		THREADS_PER_RUN*sizeof(CUDA_Seed), cudaMemcpyHostToDevice, plan->stream));

	dim3 dimBlock(512, 1);
	dim3 dimGrid(THREADS_PER_RUN/512, 1);
	TrackFibers_kernel<<<dimGrid, dimBlock, 0, plan->stream>>>
		(plan->d_pSamplePoints, plan->d_pSampleWeights, plan->d_pSeeds, 
		 plan->d_pFibers, in_stepSize, in_maxLength, in_faThreshold, 
		 in_angleThreshold, in_MLSSize, in_innerSteps, num);


	cutilSafeCall(cudaGetLastError());

	cutilSafeCall(cudaMemcpyAsync(plan->h_pFibers, plan->d_pFibers, 
		THREADS_PER_RUN*sizeof(CUDA_Fiber), cudaMemcpyDeviceToHost, plan->stream));
}

extern "C"
void CUDA_BindTextures(CDTI *DTI, CScalar *MASK, GPUPlan *plans)
{
	int w = DTI->GetWidth();
	int h = DTI->GetHeight();
	int d = DTI->GetDepth();
	int n = DTI->GetN();

	float *dti = DTI->GetDTIData();
	float *mask = MASK->GetScalarData();

	/* create textures and bind them */

	/* MASK */
	cudaExtent extent = make_cudaExtent(w, h, d);
	cudaChannelFormatDesc desc = cudaCreateChannelDesc<float>();
	cutilSafeCall(cudaMalloc3DArray(&(plans->d_pMASK), &desc, extent));

	cudaMemcpy3DParms copyParams = {0};
	copyParams.srcPtr   = make_cudaPitchedPtr((void*)(mask), w*sizeof(float), w, h);
	copyParams.dstArray = plans->d_pMASK;
	copyParams.extent   = extent;
	copyParams.kind     = cudaMemcpyHostToDevice;
	cutilSafeCall(cudaMemcpy3D(&copyParams));

	texMASK.normalized = false;                    
	texMASK.filterMode = cudaFilterModeLinear;      
	texMASK.addressMode[0] = cudaAddressModeClamp;
	texMASK.addressMode[1] = cudaAddressModeClamp;
	texMASK.addressMode[2] = cudaAddressModeClamp;
	cutilSafeCall(cudaBindTextureToArray(texMASK, plans->d_pMASK, desc));


	/* DTI */
	for (int k = 0; k < 2; ++k){
		desc = cudaCreateChannelDesc<float4>();
		cutilSafeCall(cudaMalloc3DArray(&(plans->d_pDTI[k]), &desc, extent));

		float *p = new float[w*h*d*4];
		for (int j = 0; j < w*h*d; ++j) {
			p[j*4+0] = dti[j*n+k*4+0];
			p[j*4+1] = dti[j*n+k*4+1];
			p[j*4+2] = dti[j*n+k*4+2];
			if (k == 0)
				p[j*4+3] = dti[j*n+k*4+3];
			else
				p[j*4+3] = 0.0f;
		}

		cudaMemcpy3DParms copyParams = {0};
		copyParams.srcPtr   = make_cudaPitchedPtr((void*)(p), w*sizeof(float4), w, h);
		copyParams.dstArray = plans->d_pDTI[k];
		copyParams.extent   = extent;
		copyParams.kind     = cudaMemcpyHostToDevice;
		cutilSafeCall(cudaMemcpy3D(&copyParams));

		if (k == 0) {
			texDTI_1.normalized = false;
			texDTI_1.filterMode = cudaFilterModeLinear;      
			texDTI_1.addressMode[0] = cudaAddressModeClamp;
			texDTI_1.addressMode[1] = cudaAddressModeClamp;
			texDTI_1.addressMode[2] = cudaAddressModeClamp;
			cutilSafeCall(cudaBindTextureToArray(texDTI_1, plans->d_pDTI[k], desc));
		} else {
			texDTI_2.normalized = false;
			texDTI_2.filterMode = cudaFilterModeLinear;      
			texDTI_2.addressMode[0] = cudaAddressModeClamp;
			texDTI_2.addressMode[1] = cudaAddressModeClamp;
			texDTI_2.addressMode[2] = cudaAddressModeClamp;
			cutilSafeCall(cudaBindTextureToArray(texDTI_2, plans->d_pDTI[k], desc));
		}

		delete[] p;
	}
}
